import torch
from torch import nn
from Func import validate_model, data_preproc_load
from timm import create_model

device = "mps"
model_name = "tf_efficientnetv2_b2"
num_classes = 4
classes = ['expressionism', 'impressionism', 'realism', 'surrealism']  # 4

model = create_model(model_name, pretrained=True)
model.load_state_dict(torch.load('Trained_Models/tf_efficientnetv2_b2_CL:4_VA:67.23_TA:94.91_23-12-05_14:02.pth',
                                 map_location=device))
model = model.to(device)

criterion = nn.CrossEntropyLoss()
criterion = criterion.to(device)
_, test_loader = data_preproc_load(data_improve=None, b_s=48, train_path='Data/train', test_path='Data/Evaluate')

curr_val_acc, avg_loss = validate_model(model, criterion, test_loader, device, num_classes, classes)
